package tuyou.kafka.consumer;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.util.Arrays;
import java.util.Collection;
import java.util.Properties;
import java.util.Scanner;

import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.common.TopicPartition;

import tuyou.kafka.producer.ProducerTest;

public class ExactlyOneConsumer {

    private static OffsetManager offsetManager = new OffsetManager("storage2");

    public static void main(String[] str) throws InterruptedException {

        System.out.println("Starting ...");
        readMessages();
    }

    private static void readMessages() throws InterruptedException {

        KafkaConsumer<String, String> consumer = createConsumer();
        consumer.subscribe(Arrays.asList(ProducerTest.topic), new MyConsumerRebalancerListener(consumer));
        processRecords(consumer);
    }

    private static KafkaConsumer<String, String> createConsumer() {
        Properties props = new Properties();
        props.put("bootstrap.servers", ProducerTest.ip);
        String consumeGroup = "cg3";

        props.put("group.id", consumeGroup);

        // Below is a key setting to turn off the auto commit.
        props.put("enable.auto.commit", "false");
        props.put("heartbeat.interval.ms", "2000");
        props.put("session.timeout.ms", "6001");

        // Control maximum data on each poll, make sure this value is bigger than the maximum single record size
        props.put("max.partition.fetch.bytes", "140");

        props.put("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        props.put("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        return new KafkaConsumer<>(props);
    }

    private static void processRecords(KafkaConsumer<String, String> consumer) {

        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(100);

            for (ConsumerRecord<String, String> record : records) {
                System.out.printf("offset = %d, key = %s, value = %s\n", record.offset(), record.key(), record.value());
                offsetManager.saveOffsetInExternalStore(record.topic(), record.partition(), record.offset());
            }
        }
    }
}

class MyConsumerRebalancerListener implements ConsumerRebalanceListener {

    private OffsetManager offsetManager = new OffsetManager("storage2");
    private Consumer<String, String> consumer;

    public MyConsumerRebalancerListener(Consumer<String, String> consumer) {
        this.consumer = consumer;
    }

    /**
     * 再均衡开始之前和消费者停止读取消息之后被调用
     * @param partitions
     */
    public void onPartitionsRevoked(Collection<TopicPartition> partitions) {

        for (TopicPartition partition : partitions) {
            offsetManager.saveOffsetInExternalStore(partition.topic(), partition.partition(), consumer.position(partition));
        }
    }

    /**
     * 重新分配分区之后和消费者开始读取消息之前
     * @param partitions
     */
    public void onPartitionsAssigned(Collection<TopicPartition> partitions) {

        for (TopicPartition partition : partitions) {
            consumer.seek(partition, offsetManager.readOffsetFromExternalStore(partition.topic(), partition.partition()));
        }
    }
}


class OffsetManager {


      private String storagePrefix;

      public OffsetManager(String storagePrefix) {
          this.storagePrefix = storagePrefix;
      }

      void saveOffsetInExternalStore(String topic, int partition, long offset) {

          try {

              FileWriter writer = new FileWriter(storageName(topic, partition), false);

              BufferedWriter bufferedWriter = new BufferedWriter(writer);
              bufferedWriter.write(offset + "");
              bufferedWriter.flush();
              bufferedWriter.close();

          } catch (Exception e) {
              e.printStackTrace();
          }
      }

      /**
       * @return he last offset + 1 for the provided topic and partition.
       */
      long readOffsetFromExternalStore(String topic, int partition) {

          try {
        	  File file = new File(storageName(topic, partition));
        	  if(!file.exists()){
        		  file.createNewFile();
        	  }

        	  Scanner scanner = new Scanner(file);
              if(scanner.hasNextLine()){
            	  return Long.parseLong(scanner.nextLine()) + 1;
              }
          } catch (Exception e) {
              e.printStackTrace();
          }

          return 0;
      }

      private String storageName(String topic, int partition) {
          return storagePrefix + "-" + topic + "-" + partition;
      }
  }
